import torch
from torchvision import models, transforms
from PIL import Image
from pathlib import Path
from sklearn.preprocessing import normalize
import numpy as np
import faiss
import shutil
import argparse
from tqdm import tqdm

# -----------------------------
# 配置模型和预处理
# -----------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.resnet50(pretrained=True)
model.eval()
model.to(device)

preprocess = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
])

def extract_feature(path):
    img = Image.open(path).convert("RGB")
    x = preprocess(img).unsqueeze(0).to(device)
    with torch.no_grad():
        feat = model(x).flatten().cpu().numpy()
    return feat

# -----------------------------
# 主函数
# -----------------------------
def main(src, dst, n_clusters=20, recursive=True):
    src = Path(src)
    dst = Path(dst)
    dst.mkdir(exist_ok=True, parents=True)

    pattern = "**/*" if recursive else "*"
    image_paths = [p for p in src.rglob(pattern) if p.is_file() and p.suffix.lower() in [".jpg",".jpeg",".png",".bmp",".gif"]]

    if not image_paths:
        print("⚠️ 没有找到图片")
        return

    print(f"找到 {len(image_paths)} 张图片，开始提取特征...")

    features = []
    for p in tqdm(image_paths):
        feat = extract_feature(p)
        features.append(feat)
    features = np.vstack(features)
    features = normalize(features, axis=1)  # L2 normalize

    print("聚类中...")
    dim = features.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(features)
    D, I = index.search(features, 1)  # 每张图片最近邻索引
    # 用 sklearn KMeans 代替 FAISS 聚类更直观
    from sklearn.cluster import KMeans
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    labels = kmeans.fit_predict(features)

    # -----------------------------
    # 移动文件
    # -----------------------------
    for path, label in zip(image_paths, labels):
        cluster_dir = dst / f"cluster_{label}"
        cluster_dir.mkdir(parents=True, exist_ok=True)
        shutil.copy(path, cluster_dir / path.name)

    print(f"✅ 整理完成，共 {n_clusters} 个簇")

# -----------------------------
# CLI
# -----------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FAISS + PyTorch 相册整理")
    parser.add_argument("--src", required=True, help="源目录")
    parser.add_argument("--dst", required=True, help="目标目录")
    parser.add_argument("--clusters", type=int, default=20, help="聚类数")
    parser.add_argument("--no-recursive", action="store_true", help="不递归扫描")
    args = parser.parse_args()

    main(args.src, args.dst, n_clusters=args.clusters, recursive=not args.no_recursive)

